import pandas as pd
import matplotlib.pyplot as plt
import base64
import numpy as np
import io
from cachetools import TTLCache
import akshare as ak
import time

pd.set_option('display.max_columns', None)  # 显示所有列
pd.set_option('display.width', None)  # 自动调整宽度（避免换行）
pd.set_option('display.max_colwidth', None)  # 显示完整列内容

# 设置全局字体（确保系统有该字体，如Windows的SimHei、Mac的PingFang SC）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 替换为你的系统支持的中文字体
plt.rcParams['axes.unicode_minus'] = False    # 解决负号显示问题

code_cache = TTLCache(maxsize=1000000, ttl=86400)  # 股票代码缓存24小时
hist_cache = TTLCache(maxsize=1000000, ttl=3600)   # 历史数据缓存1小时
#print(ak.stock_info_sz_name_code().head())

def get_gdp_data():
    """获取GDP数据"""
    gdp_df = ak.macro_china_gdp()
    print("----gdp")
    print(gdp_df.head())
    gdp_df = gdp_df[['季度', '国内生产总值-绝对值', '国内生产总值-同比增长']]
    gdp_df.columns = ['quarter', 'gdp', 'gdp_yoy']

    # 处理季度数据转换为日期
    def parse_quarter(q_str):
        # 提取年份
        year = q_str.split('年第')[0]

        # 处理季度部分
        q_part = q_str.split('年第')[1].replace('季度', '')

        if '-' in q_part:  # 处理范围季度如"1-2"
            end_q = int(q_part.split('-')[1])  # 取结束季度
        else:  # 处理单个季度如"1"
            end_q = int(q_part)

        # 转换为季度末日期
        month = end_q * 3  # 季度转换为月份(1季度=3月，2季度=6月等)
        return pd.to_datetime(f"{year}-{month:02d}-01") + pd.offsets.MonthEnd(0)

    gdp_df['quarter'] = gdp_df['quarter'].apply(parse_quarter)

    # 设置索引并按时间排序
    gdp_df.set_index('quarter', inplace=True)
    gdp_df.sort_index(inplace=True)
    return gdp_df

def get_industrial_data():
    """获取工业增加值数据"""
    ind_df = ak.macro_china_gdp()
    print("----gdp")
    print(ind_df.head())
    ind_df = ind_df[['季度', '第一产业-绝对值', '第一产业-同比增长']]
    ind_df.columns = ['quarter', 'ind', 'ind_yoy']

    # 处理季度数据转换为日期
    def parse_quarter(q_str):
        # 提取年份
        year = q_str.split('年第')[0]

        # 处理季度部分
        q_part = q_str.split('年第')[1].replace('季度', '')

        if '-' in q_part:  # 处理范围季度如"1-2"
            end_q = int(q_part.split('-')[1])  # 取结束季度
        else:  # 处理单个季度如"1"
            end_q = int(q_part)

        # 转换为季度末日期
        month = end_q * 3  # 季度转换为月份(1季度=3月，2季度=6月等)
        return pd.to_datetime(f"{year}-{month:02d}-01") + pd.offsets.MonthEnd(0)

    ind_df['quarter'] = ind_df['quarter'].apply(parse_quarter)

    # 设置索引并按时间排序
    ind_df.set_index('quarter', inplace=True)
    ind_df.sort_index(inplace=True)
    return ind_df

def get_cpi_data():
    """获取CPI数据"""
    cpi_df = ak.macro_china_cpi()
    print("---cpi")
    print(cpi_df.head())
    cpi_df = cpi_df[['月份', '全国-当月']]
    cpi_df.columns = ['month', 'cpi']

    cpi_df['month'] = cpi_df['month'].str.replace('年','-')
    cpi_df['month'] = cpi_df['month'].str.replace('月份','')
    cpi_df['month'] = pd.to_datetime(cpi_df['month'], format='%Y-%m')
    print("---cpi")
    print(cpi_df.head())
    cpi_df.sort_values('month', inplace=True)
    cpi_q_df = cpi_df.resample('Q', on='month').mean().reset_index()

    return cpi_q_df

def get_ppi_data():
    """获取PPI数据"""
    ppi_df = ak.macro_china_ppi()
    print("---ppi")
    print(ppi_df)
    ppi_df = ppi_df[['月份', '当月']]
    ppi_df.columns = ['month', 'ppi']

    ppi_df['month'] = ppi_df['month'].str.replace('年', '-')
    ppi_df['month'] = ppi_df['month'].str.replace('月份', '')
    ppi_df['month'] = pd.to_datetime(ppi_df['month'], format='%Y-%m')
    print("---cpi")
    print(ppi_df.head())
    ppi_df.sort_values('month', inplace=True)
    ppi_q_df = ppi_df.resample('Q', on='month').mean().reset_index()
    return ppi_q_df


def calculate_growth_momentum_factor():
    """
    计算经济增长动量因子
    因子公式：(GDP同比增速 + 工业增加值同比增速)/2
    """
    gdp_df = get_gdp_data()
    ind_df = get_industrial_data()

    # 合并数据
    merged_df = pd.merge(gdp_df, ind_df, left_index=True, right_index=True, how='inner')

    # 计算因子值
    merged_df['growth_momentum'] = (merged_df['gdp_yoy'] + merged_df['ind_yoy']) / 2

    # 标准化因子
    merged_df['growth_momentum_norm'] = (merged_df['growth_momentum'] - merged_df['growth_momentum'].mean()) / \
                                        merged_df['growth_momentum'].std()

    return merged_df[['gdp_yoy', 'ind_yoy', 'growth_momentum', 'growth_momentum_norm']]


def calculate_inflation_spread_factor():
    """
    计算通胀利差因子
    因子公式：CPI - PPI (反映上下游价格传导情况)
    """
    cpi_df = get_cpi_data()
    ppi_df = get_ppi_data()

    # 合并数据
    merged_df = pd.merge(cpi_df, ppi_df, left_index=True, right_index=True, how='inner')

    # 计算因子值
    merged_df['inflation_spread'] = merged_df['cpi'] - merged_df['ppi']

    # 标准化因子
    merged_df['inflation_spread_norm'] = (merged_df['inflation_spread'] - merged_df['inflation_spread'].mean()) / \
                                         merged_df['inflation_spread'].std()

    return merged_df[['cpi', 'ppi', 'inflation_spread', 'inflation_spread_norm']]


def fetch_single_stock(code):
    #from app.controllers.predict_controller import cache
    """使用 AKShare 获取沪深股票列表（替代新浪财经爬虫）"""
    try:
        start_date = '2024-10-10'
        period = 'daily'
        #print(end_date)
        all_stocks = []
    except Exception as e:
        print(f"失败: {str(e)}")

    time.sleep(0.1)
    try:
        # 获取单只股票历史数据
        #print(code)
        #
        cache_key_hist = f"stock_hist_{code}"
        cached_data = hist_cache.get(cache_key_hist)
        if cached_data is not None:
            print(f"股票 {code} 数据从缓存读取")
            stock_df = cached_data
        else:
            print(f"从接口获取股票 {code} 历史数据...")
            stock_df = ak.stock_zh_a_hist(
                symbol=code,
                period=period,
                start_date=start_date
            )
            hist_cache[cache_key_hist] = stock_df  # 缓存历史数据
            print(f"股票 {code} 数据已缓存")
            print(stock_df.head())

        # 数据清洗
        stock_df.rename(
            columns={
                '日期': '交易日期',
                '开盘': '开盘价',
                '收盘': '收盘价',
                '最高': '最高价',
                '最低': '最低价',
                '成交量': '成交量',
                '成交额': '成交额',
                '振幅': '振幅',
                '涨跌幅': '涨跌幅',
                '涨跌额': '涨跌额',
                '换手率': '换手率'
            },
            inplace=True
        )

        # 添加股票代码列
        stock_df['股票代码'] = code

        # 选择需要的列
        stock_df = stock_df[['交易日期', '股票代码', '收盘价','最高价','最低价','成交量','成交额','振幅','涨跌幅','涨跌额','换手率']]

        all_stocks.append(stock_df.tail(100))
        #print(all_stocks)
    except Exception as e:
        print(f"股票 {code} 获取失败: {str(e)}")

    # 合并所有股票数据
    if all_stocks:
        result = pd.concat(all_stocks, ignore_index=True)
        print(f"最终合并数据，共 {len(result)} 条记录")
        return result
    return pd.DataFrame()


def calculate_max_drawdown(prices):
    """
    计算价格序列的最大回撤（优化版）
    :param prices: 价格序列
    :return: 最大回撤(小数形式)，回撤开始索引，回撤结束索引
    """
    if len(prices) < 2:
        return 0.0, 0, 0

    # 转换为numpy数组提高性能
    if isinstance(prices, (pd.Series, pd.DataFrame)):
        prices = prices.values

    # 计算累计最大值（历史峰值）
    peak = np.maximum.accumulate(prices)

    # 避免除以零
    peak[peak == 0] = np.nan

    # 计算回撤
    drawdown = (peak - prices) / peak

    # 找出最大回撤
    max_dd = np.nanmax(drawdown)
    end_idx = np.nanargmax(drawdown)
    start_idx = np.argmax(prices[:end_idx]) if end_idx > 0 else 0

    return max_dd if not np.isnan(max_dd) else 0.0, start_idx, end_idx

def calculate_factors(stock_data: pd.DataFrame) -> pd.DataFrame:
    """
    计算因子数据
    :param stock_data: 包含多只股票历史数据的DataFrame
    :return: 包含因子数据的DataFrame
    """
    if stock_data.empty:
        return pd.DataFrame()

    # 按股票代码分组计算
    grouped = stock_data.groupby('股票代码')

    factors_list = []
    for code, group in grouped:
        group = group.sort_values('交易日期')

        # 计算收益率（每日）
        group['daily_return'] = group['收盘价'].pct_change()

        # 计算累计收益率
        group['cum_return'] = (1 + group['daily_return']).cumprod() - 1

        # 计算动量因子(20日收益率)
        group['momentum_20'] = group['收盘价'].pct_change(20)

        # 计算波动率因子(20日收益率标准差)
        returns = group['收盘价'].pct_change()
        group['volatility_20'] = returns.rolling(20).std()

        # 计算成交量因子(10日成交量均值的偏离度)
        volume_ma = group['成交量'].rolling(10).mean()
        group['volume_alpha'] = (group['成交量'] - volume_ma) / volume_ma * group['收盘价'].pct_change()

        # 计算最大回撤
        max_dd, start_idx, end_idx = calculate_max_drawdown(group['收盘价'])
        group['max_drawdown'] = max_dd  # 整个序列的最大回撤

        # 计算滚动最大回撤（过去30日）
        group['rolling_30d_drawdown'] = group['收盘价'].rolling(30).apply(
            lambda x: calculate_max_drawdown(x)[0], raw=True
        )

        factors_list.append(group)

    return pd.concat(factors_list)


def plot_factors(factor_data: pd.DataFrame, stock_code: str, growth_factor, inflation_factor):
    """
    绘制指定股票的因子时序图
    :param factor_data: 包含因子数据的DataFrame
    :param stock_code: 要绘制的股票代码
    """
    stock_data = factor_data[factor_data['股票代码'] == stock_code]
    if stock_data.empty:
        print(f"未找到股票 {stock_code} 的数据")
        return

    plt.figure(figsize=(14, 18))

    # 1. 价格和累计收益率
    plt.subplot(7, 1, 1)
    plt.plot(stock_data['交易日期'], stock_data['收盘价'], label='收盘价', color='blue')
    plt.title(f'价格走势')
    plt.legend(loc='upper left')
    plt.grid(True)

    ax2 = plt.gca().twinx()
    ax2.plot(stock_data['交易日期'], stock_data['cum_return'] * 100,
             label='累计收益率(%)', color='green', linestyle='--')
    ax2.legend(loc='upper right')

    # 动量因子
    plt.subplot(7, 1, 2)
    plt.plot(stock_data['交易日期'], stock_data['momentum_20'], label='20日动量')
    plt.title(f'动量因子')
    plt.legend()
    plt.grid(True)

    # 波动率因子
    plt.subplot(7, 1, 3)
    plt.plot(stock_data['交易日期'], stock_data['volatility_20'], label='20日波动率', color='orange')
    plt.title(f'波动率因子')
    plt.legend()
    plt.grid(True)

    # 成交量因子
    plt.subplot(7, 1, 4)
    plt.plot(stock_data['交易日期'], stock_data['volume_alpha'], label='成交量Alpha', color='green')
    plt.title(f'成交量因子')
    plt.legend()
    plt.grid(True)

    # 4. 最大回撤
    plt.subplot(7, 1, 5)
    plt.plot(stock_data['交易日期'], stock_data['rolling_30d_drawdown'] * 100,
             label='滚动30日最大回撤(%)', color='brown')
    plt.axhline(y=stock_data['max_drawdown'].iloc[0] * 100, color='black',
                linestyle='--', label='历史最大回撤')
    plt.title(f'最大回撤分析')
    plt.legend()
    plt.grid(True)

    # 经济增长动量因子
    plt.subplot(7, 1, 6)
    plt.plot(growth_factor.index, growth_factor['growth_momentum'], label='经济增长动量因子', color='blue')
    plt.fill_between(growth_factor.index, growth_factor['growth_momentum'], alpha=0.1, color='blue')
    plt.title('经济增长动量因子 (GDP增速与工业增加值增速的平均)')
    plt.ylabel('因子值(%)')
    plt.legend()
    plt.grid(True)

    # 通胀利差因子
    plt.subplot(7, 1, 7)
    plt.plot(inflation_factor.index, inflation_factor['inflation_spread'], label='通胀利差因子(CPI-PPI)', color='red')
    plt.fill_between(inflation_factor.index, inflation_factor['inflation_spread'], alpha=0.1, color='red')
    plt.title('通胀利差因子 (CPI与PPI差值)')
    plt.ylabel('因子值(%)')
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', dpi=100)
    plt.close()
    buf.seek(0)
    return base64.b64encode(buf.read()).decode('utf-8')